﻿using Ndns.Packet.Abstractions;
using Ndns.Packet.Parser;
using Ndns.ServerHost.Abstractions;
using System;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;

namespace Ndns.Net
{
    public class DnsTcpServer : IDnsServer, IDisposable
    {

        public DnsProtocol Protocol => DnsProtocol.Tcp;

        public EndPoint LocalEndPoint { get; set; }

        private volatile int _state;
        public DnsServerState State => (DnsServerState)_state;

        public object Tag { get; set; }

        public IDnsRequestHandler Handler { get; private set; }

        private bool _enableCompression = true;
        public bool EnableCompression
        {
            get => _enableCompression;
            set
            {
                if (_state == (int)DnsServerState.Running)
                    throw new InvalidOperationException();
                _enableCompression = value;
            }
        }

        private int _backlog = 10;
        public int Backlog
        {
            get => _backlog;
            set
            {
                if (_state == (int)DnsServerState.Running)
                    throw new InvalidOperationException();
                _backlog = value;
            }
        }

        private int _readTimeout = 20000;
        public int ReadTimeout
        {
            get => _readTimeout;
            set
            {
                if (_state == (int)DnsServerState.Running)
                    throw new InvalidOperationException();
                _readTimeout = value;
            }
        }

        private int _sendTimeout = 20000;
        public int SendTimeout
        {
            get => _sendTimeout;
            set
            {
                if (_state == (int)DnsServerState.Running)
                    throw new InvalidOperationException();
                _sendTimeout = value;
            }
        }

        private bool _continuationEnabled = true;
        public bool ContinuationEnabled
        {
            get => _continuationEnabled;
            set
            {
                if (_state == (int)DnsServerState.Running)
                    throw new InvalidOperationException();
                _continuationEnabled = value;
            }
        }

        private int _continuationTimeout = 20000;
        public int ContinuationTimeout
        {
            get => _continuationTimeout;
            set
            {
                if (_state == (int)DnsServerState.Running)
                    throw new InvalidOperationException();
                _continuationTimeout = value;
            }
        }


        public void SetHandler(IDnsRequestHandler handler)
        {
            if (handler == null)
                throw new ArgumentNullException(nameof(handler));
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsTcpServer));
            if (_state == (int)DnsServerState.Running)
                throw new InvalidOperationException();

            Handler = handler;
        }

        private ReaderWriterLockSlim _rwLock = new ReaderWriterLockSlim();
        private CancellationTokenSource _cancelClients;
        private volatile int _version;
        private Socket _socket;
        
        private DnsPacketParser _packetParser;

        public DnsTcpServer()
        {
            _packetParser = new DnsPacketParser();
        }
        public DnsTcpServer(DnsPacketParser packetParser)
        {
            _packetParser = packetParser;
        }
        public DnsTcpServer(IPEndPoint localEP)
        {
            LocalEndPoint = localEP ?? throw new ArgumentNullException(nameof(localEP));
            _packetParser = new DnsPacketParser();
        }
        public DnsTcpServer(IPEndPoint localEP, DnsPacketParser packetParser)
        {
            LocalEndPoint = localEP ?? throw new ArgumentNullException(nameof(localEP));
            _packetParser = packetParser;
        }

        public void Start()
        {
            //检查前置条件
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsTcpServer));
            if (LocalEndPoint == null)
                throw new InvalidOperationException();
            if (Handler == null)
                throw new InvalidOperationException();
            if (_state == (int)DnsServerState.Running)
                throw new InvalidOperationException();

            //通知即将启动
            try
            {
                Handler.OnStarting(this);
            }
            catch (Exception)
            {
                _state = (int)DnsServerState.Error;
                return;
            }

            //新建Socket
            _socket = new Socket(LocalEndPoint.AddressFamily, SocketType.Stream, ProtocolType.Tcp);

            try
            {
                //绑定到端口
                _socket.Bind(LocalEndPoint);
                _socket.Listen(_backlog);
            }
            catch (SocketException ex)
            {
                //设置状态并通知错误
                _state = (int)DnsServerState.Error;
                _socket.Dispose();
                try
                {
                    Handler.OnFatalError(this, ex);
                }
                catch (Exception)
                {
                    // Do nothing
                }
                return;
            }

            //设置状态并通知启动
            _cancelClients = new CancellationTokenSource();
            _state = (int)DnsServerState.Running;
            _version = unchecked(_version + 1);
            try
            {
                Handler.OnStarted(this);
            }
            catch (Exception)
            {
                // Do nothing
            }

            //开始接收
            Task.Factory.StartNew(SocketWorkerLoop, this);
        }

        public void Stop()
        {
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsTcpServer));

            if (_state == (int)DnsServerState.Running)
            {
                try
                {
                    Handler.OnStopping(this);
                }
                catch (Exception)
                {
                    // Do nothing
                }

                _version = unchecked(_version + 1);

                _cancelClients.Cancel();

                _rwLock.EnterWriteLock();
                _socket.Dispose();
                _socket = null;
                _cancelClients = null;
                _rwLock.ExitWriteLock();

            }
            _state = (int)DnsServerState.Stopped;

            try
            {
                Handler.OnStopped(this);
            }
            catch (Exception)
            {
                // Do nothing
            }
        }

        private static async Task SocketWorkerLoop(object state)
        {
            var server = (DnsTcpServer)state;
            var version = server._version;

            while (version == server._version)
            {
                Socket client = null;
                Task<Socket> taskAccept = null;
                server._rwLock.EnterReadLock();
                try
                {
                    taskAccept = server._socket.AcceptAsync();
                }
                catch (ObjectDisposedException)
                {
                    break;
                }
                finally
                {
                    server._rwLock.ExitReadLock();
                }

                try
                {
                    client = await taskAccept.ConfigureAwait(false);
                }
                catch (ObjectDisposedException)
                {
                    break;
                }
                catch (Exception ex)
                {
                    Interlocked.CompareExchange(ref server._state, (int)DnsServerState.Error, (int)DnsServerState.Running);
                    server.Handler.OnFatalError(server, ex);
                    break;
                }

                HandleConnection(server, version, client);
            }
        }

        private static void HandleConnection(DnsTcpServer server, int version, Socket client)
        {
            Task.Run(() => HandleConnectionAsync(server, version, client, false));
        }

        private static async Task HandleConnectionAsync(DnsTcpServer server, int version, Socket client, bool isContinuation)
        {
            CancellationToken globalCancelToken = default(CancellationToken);

            //获得各个参数
            var readTimeout = server._readTimeout;
            var sendTimeout = server._sendTimeout;
            if (isContinuation)
                readTimeout = server._continuationTimeout;
            var continuationEnabled = server._continuationEnabled;
            var enableCompression = server._enableCompression;

            //获取全局取消通知
            server._rwLock.EnterReadLock();
            if (version == server._version)
            {
                globalCancelToken = server._cancelClients.Token;
            }
            else
            {
                server._rwLock.ExitReadLock();
                return;
            }
            server._rwLock.ExitReadLock();
            
            using (var ns = new NetworkStream(client, false))
            using (var ms = new MemoryStream(512))
            {
                DnsPacket packet = null;

                //设置超时
                CancellationToken readCancel = default(CancellationToken);
                if (readTimeout > 0)
                    readCancel = (new CancellationTokenSource(readTimeout)).Token;
                readCancel = CancellationTokenSource.CreateLinkedTokenSource(globalCancelToken, readCancel).Token;

                //从连接流中读取
                using (readCancel.Register(client.Dispose))
                {
                    try
                    {
                        await ns.CopyBytesToAsync(ms, 2).ConfigureAwait(false);
                        ms.Seek(0, SeekOrigin.Begin);
                        int expectedLength = ms.ReadNetworkUInt16();
                        await ns.CopyBytesToAsync(ms, expectedLength).ConfigureAwait(false);
                    }
                    catch (Exception)
                    {
                        client.Shutdown(SocketShutdown.Both);
                        client.Dispose();
                        return;
                    }
                }

                //解析数据包
                try
                {
                    ms.Seek(2, SeekOrigin.Begin);
                    server._packetParser.TryParse(ms, out packet, out var c);
                }
                catch (Exception)
                {
                    // Do nothing
                }

                //检查解析结果
                if (packet == null)
                {
                    client.Shutdown(SocketShutdown.Both);
                    client.Dispose();
                    return;
                }
                
                //上锁防止服务器被停止
                Task<DnsPacket> taskResponse = null;
                server._rwLock.EnterReadLock();
                try
                {
                    //调用回调函数
                    if (version == server._version)
                    {
                        taskResponse = server.Handler.OnReceivedAsync(server, packet, client.RemoteEndPoint, globalCancelToken);
                    }
                }
                catch (Exception)
                {
                    // Do nothing
                }
                finally
                {
                    //解锁
                    server._rwLock.ExitReadLock();
                }

                //等待回调函数结束
                DnsPacket response = null;
                try
                {
                    response = await taskResponse.ConfigureAwait(false);
                }
                catch (Exception)
                {
                    // Do nothing
                }

                //检查回调函数结果
                if (response == null)
                {
                    client.Shutdown(SocketShutdown.Both);
                    client.Dispose();
                    return;
                }
                ms.Seek(2, SeekOrigin.Begin);
                ms.SetLength(2);
                int count = enableCompression ? response.CopyCompressedTo(ms) : response.CopyTo(ms);
                ms.Seek(0, SeekOrigin.Begin);
                ms.WriteNetworkUInt16((ushort)count);
                ms.Seek(0, SeekOrigin.Begin);
                
                //准备发送结果给客户端
                server._rwLock.EnterReadLock();
                Task taskSend = null;
                if (version == server._version)
                {
                    taskSend = ms.CopyToAsync(ns);
                }
                server._rwLock.ExitReadLock();

                //设置超时
                CancellationToken sendCancel = default(CancellationToken);
                if (sendTimeout > 0)
                    sendCancel = (new CancellationTokenSource(sendTimeout)).Token;
                sendCancel = CancellationTokenSource.CreateLinkedTokenSource(globalCancelToken, sendCancel).Token;

                //等待发送完毕
                using (sendCancel.Register(client.Dispose))
                {
                    try
                    {
                        if (taskSend != null)
                            await taskSend.ConfigureAwait(false);
                    }
                    catch (Exception)
                    {
                        client.Shutdown(SocketShutdown.Both);
                        client.Dispose();
                        return;
                    }
                }

            }

            //继续等待下一个请求
            Task dummy = null;
            if (continuationEnabled)
            {
                dummy = Task.Run(() => HandleConnectionAsync(server, version, client, true));
            }
            else
            {
                client.Shutdown(SocketShutdown.Both);
                client.Dispose();
            }
        }

        #region IDisposable Support
        private bool disposedValue = false; // 要检测冗余调用

        protected virtual void Dispose(bool disposing)
        {
            if (disposing)
            {
                // 释放托管状态(托管对象)。
                _socket?.Dispose();
            }

            // 释放未托管的资源(未托管的对象)并在以下内容中替代终结器。
            // 将大型字段设置为 null。
            _socket = null;
            _state = (int)DnsServerState.Stopped;

            disposedValue = true;
        }

        // 添加此代码以正确实现可处置模式。
        public void Dispose()
        {
            // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
            Dispose(true);
        }
        #endregion
    }
}
